import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 读取数据
amount = pd.read_csv('Amount.csv')
close = pd.read_csv('Close.csv')
index = pd.read_csv('index.csv')

# 数据预处理
# 将第一列设为日期索引
amount.rename(columns={'Unnamed: 0': 'date'}, inplace=True)
close.rename(columns={'Unnamed: 0': 'date'}, inplace=True)
index.rename(columns={'Unnamed: 0': 'date'}, inplace=True)

amount['date'] = pd.to_datetime(amount['date'])
close['date'] = pd.to_datetime(close['date'])
index['date'] = pd.to_datetime(index['date'])

amount.set_index('date', inplace=True)
close.set_index('date', inplace=True)
index.set_index('date', inplace=True)

# 按照测试要求，将样本等分为十个组别
def backtest_strategy():
    """
    对股票数据进行回测分析，将股票分为10组，每组构建基于成交额的股票组合，计算其收益表现
    
    该函数实现以下功能：
    1. 将所有股票等分为10组
    2. 每组每日选取成交额前300只股票构建等权投资组合
    3. 计算每日组合收益、换手率和相对于等权指数的超额收益(Alpha)
    4. 累计每日Alpha得到累计超额收益曲线
    
    Returns:
        list: 包含10个字典的列表，每个字典代表一组的回测结果，包含以下键值：
            - 'group': 组别编号 (1-10)
            - 'dates': 交易日期列表
            - 'returns': 每日收益率列表
            - 'alphas': 每日Alpha列表
            - 'cumulative_alphas': 累计Alpha列表
            - 'turnovers': 每日换手率列表
    """
    # 用于存储每组的结果
    group_results = []
    
    # 对每组进行回测 (10组)
    for group in range(10):
        print(f"正在处理第 {group + 1} 组...")
        
        # 存储每日收益率和alpha
        daily_returns = []
        daily_alphas = []
        daily_turnovers = []
        dates = []
        
        # 初始资金
        total_capital = 300.0  # 300份，每份1元
        capital_per_stock = 1.0
        
        # 持仓股票列表
        current_positions = []
        current_position_values = []
        
        # 每日循环
        for i in range(len(close)):
            current_date = close.index[i]
            dates.append(current_date)
            
            # 获取当日的收盘价和成交额数据
            close_data = close.loc[current_date]
            amount_data = amount.loc[current_date]
            
            # 确定当前组别的股票范围 (每组约324只股票)
            total_stocks = len(close.columns)
            stocks_per_group = total_stocks // 10
            start_idx = group * stocks_per_group
            end_idx = start_idx + stocks_per_group if group < 9 else total_stocks
            group_stocks = close.columns[start_idx:end_idx]
            
            # 在该组内选择成交额最大的300只股票（如果该组股票数不足300则选所有）
            group_amount_data = amount_data[group_stocks]
            top_stocks = group_amount_data.nlargest(min(300, len(group_stocks))).index.tolist()
            
            # 如果是第一天，建立初始仓位
            if i == 0:
                current_positions = top_stocks[:min(300, len(top_stocks))]
                current_position_values = [capital_per_stock] * len(current_positions)
            else:
                # 计算前一天的持仓收益
                prev_close_data = close.loc[close.index[i-1]]
                today_returns = []
                
                # 计算每只股票的收益
                for j, stock in enumerate(current_positions):
                    if stock in close_data and stock in prev_close_data and prev_close_data[stock] > 0:
                        stock_return = close_data[stock] / prev_close_data[stock] - 1
                        current_position_values[j] *= (1 + stock_return)
                        today_returns.append(stock_return)
                    else:
                        today_returns.append(0)
                
                # 确定需要卖出和继续持有的股票
                stocks_to_sell = []
                stocks_to_keep = []
                keep_values = []
                
                for j, stock in enumerate(current_positions):
                    if stock in top_stocks:
                        # 继续持有该股票
                        stocks_to_keep.append(stock)
                        keep_values.append(current_position_values[j])
                    else:
                        # 卖出该股票
                        stocks_to_sell.append((stock, current_position_values[j]))
                
                # 确定需要买入的股票（在目标股票列表中但当前未持有的股票）
                stocks_to_buy = [s for s in top_stocks if s not in current_positions]
                
                # 计算换手率：新买入股票数占当前持仓股票数的比例
                turnover = len(stocks_to_buy) / len(current_positions) if len(current_positions) > 0 else 0
                daily_turnovers.append(turnover)
                
                # 计算卖出股票获得的资金
                sold_capital = sum([value for stock, value in stocks_to_sell])
                
                # 更新持仓：仅保留需要继续持有的股票
                current_positions = stocks_to_keep
                current_position_values = keep_values
                
                # 如果有新股票需要买入，将卖出股票的资金平均分配买入
                if stocks_to_buy:
                    capital_per_new_stock = sold_capital / len(stocks_to_buy) if stocks_to_buy else 0
                    for stock in stocks_to_buy:
                        current_positions.append(stock)
                        current_position_values.append(capital_per_new_stock)
                
                # 计算组合总收益
                total_value = sum(current_position_values)
                portfolio_return = total_value / total_capital - 1 if total_capital > 0 else 0
                total_capital = total_value
                
                # 计算等权指数收益
                index_return = index.loc[current_date, '0']  # 等权指数收益率
                
                # 计算超额收益 (alpha)
                alpha = portfolio_return - index_return
                
                daily_returns.append(portfolio_return)
                daily_alphas.append(alpha)
        
        # 累计Alpha
        cumulative_alphas = np.cumsum(daily_alphas)
        
        # 存储结果
        group_results.append({
            'group': group + 1,
            'dates': dates[1:],  # 从第二天开始
            'returns': daily_returns,
            'alphas': daily_alphas,
            'cumulative_alphas': cumulative_alphas,
            'turnovers': daily_turnovers
        })
        
        print(f"第 {group + 1} 组完成。平均换手率: {np.mean(daily_turnovers):.4f}")
    
    return group_results

# 运行回测
print("开始运行回测...")
results = backtest_strategy()

# 绘制结果
plt.figure(figsize=(15, 8))

# 绘制每组的累计Alpha曲线
for result in results:
    plt.plot(result['dates'], result['cumulative_alphas'], label=f'第{result["group"]}组')

plt.xlabel('日期')
plt.ylabel('累计Alpha')
plt.title('十组股票的累计Alpha曲线')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig('cumulative_alpha.png', dpi=300, bbox_inches='tight')
plt.show()

# 输出统计数据
print("\n回测结果统计:")
for result in results:
    print(f"第 {result['group']} 组:")
    print(f"  平均日收益率: {np.mean(result['returns']):.4f}")
    print(f"  平均Alpha: {np.mean(result['alphas']):.4f}")
    print(f"  最终累计Alpha: {result['cumulative_alphas'][-1]:.4f}")
    print(f"  平均换手率: {np.mean(result['turnovers']):.4f}")
    print()